﻿using System;
using System.IO;
using System.Windows;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Runtime.InteropServices;
using DotNetImage.Util;
using System.Collections.Generic;

namespace DotNetImage.Imaging
{
    /// <summary>
    /// <remarks>To check the how the little-endian affects the code of conversion, GetPixel / SetPixel etc.</remarks>
    /// </summary>
    public class DotNetImage
    {
        #region Properties
        /// <summary>
        /// Width of the picture in pixels
        /// </summary>
        public int Width { get; internal set; }

        /// <summary>
        /// Height of the picture in pixels
        /// </summary>
        public int Height { get; internal set; }

        /// <summary>
        /// Number of bytes in one pixel. Generally speaking, if it's 1 or 2, then
        /// image is a grayscale one. When it's 3, 6 or 12, pixel is a color one in BGR format, 
        /// having byte, ushort or uint as a color channel and if it's 4, 8 or 16, then it's a
        /// color in format BGRA, having byte, ushort or uint as a color channel.
        /// <seealso cref="ProcessPixelFormat"/>
        /// </summary>
        public int PixelSize { get; private set; }

        /// <summary>
        /// Color channel size. As for today can be 1 (byte), 2 (ushort) and 4 (uint).
        /// </summary>
        public int ColorChannelSize { get; private set; }

        /// <summary>
        /// Number of color channels: 1 (gray), 3 (RGB) or 4 (RGBA)
        /// </summary>
        public int ColorChannelCount { get; private set; }

        /// <summary>
        /// True, when pixel is monochrome (pixel size equals to color channel size)
        /// and false otherwise.
        /// </summary>
        public bool IsMonochrome { get; private set; }

        /// <summary>
        /// Type of the pixel: PixelMonochrome, Pixel RGB, PixelRGBA
        /// </summary>
        public Type PixelType { get; private set; }

        /// <summary>
        /// Type of the color channel 
        /// </summary>
        public Type ChannelType { get; private set; }

        internal BitmapSource BitmapSource { get { return _bmp; } }
        protected BitmapSource _bmp = null;

        /// <summary>
        /// The bitmap as array of bytes. Length depends on <seealso cref="InternalPixelFormat"/>.
        /// This array will be changed by image filters.
        /// </summary>
        internal byte[] PixelData = null;

        /// <summary>
        /// The copy of the pixel data. Used in filters when original values are required but <seealso cref="PixelData"/> 
        /// can be changed by the filter.
        /// </summary>
        internal byte[] PreservedPixelData = null;

        /// <summary>
        /// Just for information if someone wants to know
        /// </summary>
        public PixelFormat InternalPixelFormat { get; private set; }

        /// <summary>
        /// Static storage to keep all pixel formats supported internally
        /// </summary>
        internal static Dictionary<PixelFormat, bool> InternalPixelFormats = new Dictionary<PixelFormat, bool>();


        private uint _ColorChannelMaxValue;
        public uint ColorChannelMaxValue { get { return _ColorChannelMaxValue; } }
        #endregion

        #region Constructors
        static DotNetImage()
        {
            InternalPixelFormats.Add(PixelFormats.Rgb24, true);
            InternalPixelFormats.Add(PixelFormats.Rgb48, true);
            InternalPixelFormats.Add(PixelFormats.Rgba64, true);
            InternalPixelFormats.Add(PixelFormats.Rgba128Float, true);
            InternalPixelFormats.Add(PixelFormats.Bgra32, true);
            InternalPixelFormats.Add(PixelFormats.Gray8, true);
            InternalPixelFormats.Add(PixelFormats.Gray16, true);
            InternalPixelFormats.Add(PixelFormats.Gray32Float, true);
        }
        public DotNetImage(BitmapSource src)
        {
            _bmp = src;
            Construct();
        }
        public DotNetImage(string fileName, int width = 0, int height = 0)
        {
            _bmp = reformatImage(DotNetImage.LoadFromFile(fileName, width, height));
            Construct();
        }
        public DotNetImage(Stream stream, int width = 0, int height = 0)
        {
            _bmp = reformatImage(DotNetImage.LoadFromStream(stream, width, height));
            Construct();
        }

        private void Construct()
        {
            // Make sure that format is supported by framework
            // TODO: throw meaningful exception
            bool isSupported = InternalPixelFormats[_bmp.Format];

            Width = _bmp.PixelWidth;
            Height = _bmp.PixelHeight;
            PixelData = _bmp.CopyPixels();

            _ColorChannelMaxValue = ColorChannelSize == 1 ? byte.MaxValue : ColorChannelSize == 2 ? ushort.MaxValue : uint.MaxValue;
        }

        #endregion

        #region Helper routines
        /// <summary>
        /// Call this method to create the copy of the original pixels. Can be useful in case of the rollback.
        /// Pixels array has to be populated by this moment.
        /// </summary>
        public void CreatePixelCopy()
        {
            PreservedPixelData = new byte[PixelData.LongLength];
            Array.Copy(PixelData, PreservedPixelData, PixelData.LongLength);
        }

        /// <summary>
        /// Takes the pixel format of the original image, sets up properties
        /// of the DotNetImage and returns pixel format to be used to convert original
        /// image to.
        /// </summary>
        /// <param name="format"></param>
        /// <returns>One of the following formats: Gray8, Gray16, Gray32Float, Rgb24, Bgra32 (there is no Rgba32), Rgb48, Rgba64, Rgba128Float</returns>
        protected PixelFormat ProcessPixelFormat(PixelFormat format)
        {
            // Calculate number of color channels from resulting format
            if (format == PixelFormats.Gray8 || format == PixelFormats.Gray16 || format == PixelFormats.Gray32Float)
            {
                ColorChannelCount = 1;
            }
            else if (format == PixelFormats.Rgb24 || format == PixelFormats.Rgb48)
            {
                ColorChannelCount = 3;
            }
            else
            {
                ColorChannelCount = 4;
            }

            // Set the flag if image is monochrome
            IsMonochrome = (ColorChannelCount == 1);

            // Set number of bytes per pixel
            PixelSize = format.BitsPerPixel / 8;

            // Number of bytes per channel
            ColorChannelSize = PixelSize / ColorChannelCount;

            // The format, picture is internally presented in. Is not supposed to be changed by filters.
            InternalPixelFormat = format;

            return format;
        }

        /// <summary>
        /// Returns the format, bitmap has to be converted to.
        /// </summary>
        /// <param name="sourceFormat"></param>
        /// <returns></returns>
        internal static PixelFormat FormatConvertor(PixelFormat sourceFormat)
        {
            PixelFormat result = PixelFormats.Bgra32;
            if (sourceFormat == PixelFormats.Bgr101010
                || sourceFormat == PixelFormats.Rgb48)
            {
                result = PixelFormats.Rgb48;
            }
            else if (sourceFormat == PixelFormats.Bgr24
                || sourceFormat == PixelFormats.Bgr32
                || sourceFormat == PixelFormats.Bgr555
                || sourceFormat == PixelFormats.Bgr565
                || sourceFormat == PixelFormats.Indexed1
                || sourceFormat == PixelFormats.Indexed2
                || sourceFormat == PixelFormats.Indexed4
                || sourceFormat == PixelFormats.Indexed8
                || sourceFormat == PixelFormats.Rgb24)
            {
                result = PixelFormats.Rgb24;
            }
            else if (sourceFormat == PixelFormats.Bgra32
                || sourceFormat == PixelFormats.Pbgra32)
            {
                result = PixelFormats.Bgra32;
            }
            else if (sourceFormat == PixelFormats.Cmyk32
                || sourceFormat == PixelFormats.Prgba64
                || sourceFormat == PixelFormats.Rgba64)
            {
                result = PixelFormats.Rgba64;
            }
            else if (sourceFormat == PixelFormats.Gray16)
            {
                result = PixelFormats.Gray16;
            }
            else if (sourceFormat == PixelFormats.BlackWhite
                || sourceFormat == PixelFormats.Gray2
                || sourceFormat == PixelFormats.Gray4
                || sourceFormat == PixelFormats.Gray8)
            {
                result = PixelFormats.Gray8;
            }
            else if (sourceFormat == PixelFormats.Gray32Float)
            {
                result = PixelFormats.Gray32Float;
            }
            else if (sourceFormat == PixelFormats.Rgb128Float
                || sourceFormat == PixelFormats.Rgba128Float
                || sourceFormat == PixelFormats.Prgba128Float)
            {
                result = PixelFormats.Rgba128Float;
            }
            else //(format == PixelFormats.Default)
            {
                throw new NotImplementedException(sourceFormat.ToString());
            }
            return result;
        }

        /// <summary>
        /// Loads image from file and returns BitmapSource
        /// </summary>
        /// <param name="fileName"></param>
        /// <param name="width"></param>
        /// <param name="height"></param>
        /// <returns></returns>
        private static BitmapSource LoadFromFile(string fileName, int width = 0, int height = 0)
        {
            Uri uri = new Uri(fileName);
            BitmapImage src = new BitmapImage();
            src.CacheOption = BitmapCacheOption.OnLoad;
            src.BeginInit();
            src.UriSource = uri;
            if (width > 0) src.DecodePixelWidth = width;
            if (height > 0) src.DecodePixelHeight = height;
            src.EndInit();
            return src;
        }

        private static BitmapSource LoadFromStream(Stream stream, int width = 0, int height = 0)
        {
            BitmapImage src = new BitmapImage();
            src.CacheOption = BitmapCacheOption.OnLoad;
            src.BeginInit();
            src.StreamSource = stream;
            if (width > 0) src.DecodePixelWidth = width;
            if (height > 0) src.DecodePixelHeight = height;
            src.EndInit();
            return reformatImage(src);
        }
        #endregion

        #region Load Image (static functions)
        /// <summary>
        /// This constructor is used when it's enough to load image scaled down (for demo purposes for example).
        /// If you need proportionally scale image, leave either width or height parameter 0. Otherwise image
        /// will be stretched / skewed according to provided width and height.
        /// </summary>
        /// <param name="fileName"></param>
        /// <param name="width">0 - when no shrink planned or provided height is used to calculate proportional scaling, otherwise width in pixels</param>
        /// <param name="height">0 - when no shrink planned or provided width is used to calculate proportional scaling, otherwise height in pixels</param>
        public static DotNetImage FromFile(string fileName, int width = 0, int height = 0)
        {
            return new DotNetImage(fileName, width, height);
        }

        /// <summary>
        /// Load image from the stream. When width and or height provided - does the scaling (stretching/shrinking), otherwise loads image as is.
        /// </summary>
        /// <param name="stream"></param>
        /// <param name="width">0 - when no shrink planned or provided height is used to calculate proportional scaling, otherwise width in pixels</param>
        /// <param name="height">0 - when no shrink planned or provided width is used to calculate proportional scaling, otherwise height in pixels</param>
        /// <returns></returns>
        public static DotNetImage FromStream(Stream stream, int width = 0, int height = 0)
        {
            return new DotNetImage(stream, width, height);
        }

        /// <summary>
        /// Converting given image into one of the supported formats. <seealso cref="FormatConvertor"/>
        /// </summary>
        /// <param name="src"></param>
        /// <returns></returns>
        private static BitmapSource reformatImage(BitmapSource src)
        {
            PixelFormat format = src.Format;
            PixelFormat targetFormat = FormatConvertor(format);

            if (format == targetFormat)
            {
                return src;
            }

            ////////// Convert the BitmapSource to a new format //////////// 
            // Use the BitmapImage created above as the source for a new BitmapSource object 
            // which is set to a gray scale format using the FormatConvertedBitmap BitmapSource.                                                
            // Note: New BitmapSource does not cache. It is always pulled when required.
            FormatConvertedBitmap newFormatedBitmapSource = new FormatConvertedBitmap();

            // BitmapSource objects like FormatConvertedBitmap can only have their properties 
            // changed within a BeginInit/EndInit block.
            newFormatedBitmapSource.BeginInit();

            // Use the BitmapSource object defined above as the source for this new  
            // BitmapSource (chain the BitmapSource objects together).
            newFormatedBitmapSource.Source = src;

            // Set the new format to Bgra32.
            newFormatedBitmapSource.DestinationFormat = targetFormat;

            newFormatedBitmapSource.EndInit();

            return newFormatedBitmapSource;
        }
        #endregion

        #region Save routines
        /// <summary>
        /// Saves image to file system. Original image is not changed. All changes are applied to WritableBitmap and then saved.
        /// </summary>
        /// <param name="img"></param>
        /// <param name="fileName">Fully qualified file name</param>
        /// <param name="encoder"></param>
        internal void Save(string fileName, BitmapEncoder encoder)
        {
            WriteableBitmap wb = Finalize();
            encoder.Frames.Add(BitmapFrame.Create(wb));
            using (var filestream = new FileStream(fileName, FileMode.Create))
            {
                encoder.Save(filestream);
            }
        }

        /// <summary>
        /// Saves image to file system. Original image is not changed. All changes are applied to WritableBitmap and then saved.
        /// </summary>
        /// <param name="fileName">Fully qualified file name</param>
        /// <param name="fileFormat">One of the supported image formats</param>
        public void Save(string fileName, SupportedFileFormat fileFormat)
        {
            BitmapEncoder encoder = fileFormat.GetEncoder();
            Save(fileName, encoder);
        }


        protected WriteableBitmap Finalize()
        {
            WriteableBitmap wb = new WriteableBitmap(Width, Height, _bmp.DpiX, _bmp.DpiY, _bmp.Format, _bmp.Palette);
            wb.WritePixels(new Int32Rect(0, 0, Width, Height), PixelData, PixelSize * Width, 0);
            return wb;
        }
        #endregion

        #region Get/Set pixel
        #region GetPixel
        public Pixel GetPixel(uint x, uint y)
        {
            return GetPixel((int)x, (int)y);
        }
        public Pixel GetPixel(int x, int y)
        {
            int offset = (y * Width + x) * PixelSize;
            return GetPixel(PixelData, offset, ColorChannelSize, ColorChannelCount);
        }
        public Pixel GetPreservedPixel(int x, int y)
        {
            int offset = (y * Width + x) * PixelSize;
            Pixel result = GetPixel(PreservedPixelData, offset, ColorChannelSize, ColorChannelCount);
            result.MaxValue = ColorChannelMaxValue;
            return result;
        }
        public Pixel GetPreservedPixel(uint x, uint y)
        {
            return GetPreservedPixel((int)x, (int)y);
        }
        private static Pixel GetPixel(byte[] data, int initialOffset, int channelSize, int channelCount)
        {
            Pixel result = new Pixel();
            result.Red = Extract(data, initialOffset, ColorChannelIndex.Red, channelSize);
            if (channelCount >= 3)
            {
                result.Green = Extract(data, initialOffset, ColorChannelIndex.Green, channelSize);
                result.Blue = Extract(data, initialOffset, ColorChannelIndex.Blue, channelSize);
            }
            if (channelCount == 4)
            {
                result.Alpha = Extract(data, initialOffset, ColorChannelIndex.Alpha, channelSize);
            }
            return result;
        }

        /// <summary>
        /// 
        /// </summary>
        /// <param name="data">Represents byte array, containing pixel data</param>
        /// <param name="initialOffset">Offset from the beginning of the array to the first byte of the required pixel</param>
        /// <param name="index">Color Channel Number</param>
        /// <param name="channelSize">Size of the color channel in bytes</param>
        /// <returns></returns>
        private static uint Extract(byte[] data, int initialOffset, ColorChannelIndex index, int channelSize)
        {
            int offset = initialOffset + (int)index * channelSize;
            switch (channelSize)
            {
                case 1:
                    return data[offset];
                case 2:
                    return BitConverter.ToUInt16(data, offset);
                case 4:
                    return BitConverter.ToUInt32(data, offset);
                default:
                    throw new NotSupportedException(channelSize.ToString());
            }
        }
        #endregion

        #region SetPixel
        public void SetPixel(Pixel pix, uint x, uint y)
        {
            SetPixel (pix, (int)x, (int)y);
        }
        public void SetPixel(Pixel pix, int x, int y)
        {
            int offset = (y * Width + x) * PixelSize;
            SetPixel(PixelData, pix, ColorChannelIndex.Red, offset, ColorChannelSize);
            if (ColorChannelCount >= 3)
            {
                SetPixel(PixelData, pix, ColorChannelIndex.Green, offset, ColorChannelSize);
                SetPixel(PixelData, pix, ColorChannelIndex.Blue, offset, ColorChannelSize);
            }
            if (ColorChannelCount == 4)
            {
                SetPixel(PixelData, pix, ColorChannelIndex.Alpha, offset, ColorChannelSize);
            }
        }

        /// <summary>
        /// TODO: Check about little-endians
        /// </summary>
        /// <param name="value"></param>
        /// <param name="index"></param>
        /// <param name="channelSize"></param>
        /// <returns></returns>
        private static void SetPixel(byte[] pixelData, Pixel pix, ColorChannelIndex index, int initialOffset, int channelSize)
        {
            uint value = pix.GetChannel(index);
            int offset = initialOffset + (int)index * channelSize;
            switch (channelSize)
            {
                case 1:
                    pixelData[offset] = (byte)value;
                    return;
                case 2:
                    byte[] dataUshort = BitConverter.GetBytes((ushort)value);
                    //if (BitConverter.IsLittleEndian)
                    //{
                    //    Array.Reverse(dataUshort);
                    //}
                    Array.Copy(dataUshort, 0, pixelData, offset, 2);
                    return;
                case 4:
                    byte[] dataUint = BitConverter.GetBytes(value);
                    //if (BitConverter.IsLittleEndian)
                    //{
                    //    Array.Reverse(dataUint);
                    //}
                    Array.Copy(dataUint, 0, pixelData, offset, 4);
                    return;
                default:
                    throw new NotSupportedException(channelSize.ToString());
            }
        }
        #endregion
        #endregion
    }



    public static class BitmapSourceExtension
    {
        /// <summary>
        /// Extension method to copy bitmap source content into internal representation
        /// </summary>
        /// <param name="source"></param>
        /// <returns></returns>
        public static byte[] CopyPixels(this BitmapSource source)
        {
            int pixelSize = source.Format.BitsPerPixel / 8;
            byte[] pixels = new byte[source.PixelWidth * source.PixelHeight * pixelSize];
            int stride = source.PixelWidth * pixelSize;
            GCHandle pinnedPixels = GCHandle.Alloc(pixels, GCHandleType.Pinned);
            source.CopyPixels(
              new Int32Rect(0, 0, source.PixelWidth, source.PixelHeight),
              pinnedPixels.AddrOfPinnedObject(),
              pixels.Length,
              stride);
            pinnedPixels.Free();
            return pixels;
        }
    }
}
